# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import json

import numpy as np
import pandas as pd

from absl import (flags, app)

FLAGS = flags.FLAGS

flags.DEFINE_string('og_log_path', '', '')

flags.DEFINE_string('include_symbols', '', '')


def read_order_event(path):
  symbols = set()
  if FLAGS.include_symbols:
    symbols = set(FLAGS.include_symbols.split(','))
  with open(path, 'r') as fi:
    data = json.load(fi)
  oe = []
  for e in data:
    if e['type'] != 'ORDER_EVENT':
      continue
    if symbols and e['event']['symbol'] not in symbols:
      continue
    e['event']['record_time'] = e['record_timestamp']
    oe.append(e['event'])
  oe = pd.DataFrame(oe)
  oe['event_time'] = oe['event_time'].astype(int)
  return oe


def output_percentile(ts, title):
  print('==================')
  print(title)
  print(ts.describe())
  for i in np.arange(0, 101, 5):
    print('%s-percentile: %f' % (i, np.percentile(ts, i)))


def main(argv):
  og_log_path = FLAGS.og_log_path
  assert og_log_path
  oe = read_order_event(og_log_path)
  for event_type, sub_df in oe.groupby('type'):
    print(event_type, sub_df.shape[0])
  print('')

  rejected_count = (oe['type'] == 'ORDER_REJECTED').sum()
  print('Order submit reject rate: %f\n' % (rejected_count / oe.shape[0]))

  # Latency
  same_order_threshold = 20e9
  submit_latency = []
  cancel_latency = []
  cancel_submit_latency = []
  last_event_by_order = dict()
  last_event_by_type = dict()
  for _, row in oe.iterrows():
    iid = row['internal_order_id']
    event_time = row['record_time']
    event_type = row['type']
    if iid not in last_event_by_order:
      last_event_by_order[iid] = dict()
    if event_type == 'ORDER_ACCEPTED':
      submitted_ts = last_event_by_order[iid].get('ORDER_SUBMITTED', 0)
      submit_lat = event_time - submitted_ts
      if submit_lat < same_order_threshold:
        submit_latency.append(submit_lat)
    elif event_type in ['CANCEL_CONFIRMED', 'CANCEL_ERROR']:
      cancel_ts = last_event_by_order[iid].get('CANCEL_SUBMITTED', 0)
      cancel_lat = event_time - cancel_ts
      if cancel_lat < same_order_threshold:
        cancel_latency.append(cancel_lat)
    elif event_type == 'ORDER_SUBMITTED':
      last_cancel_submit_event = last_event_by_type.get('CANCEL_CONFIRMED', None)
      if last_cancel_submit_event is not None:
        cancel_submit_latency.append(event_time - last_cancel_submit_event['event_time'])
    last_event_by_order[iid][event_type] = event_time
    last_event_by_type[event_type] = row

  output_percentile(pd.Series(submit_latency) * 1e-9, 'Submit Latency')
  print('')
  output_percentile(pd.Series(cancel_latency) * 1e-9, 'Cancel Latency')
  print('')
  output_percentile(pd.Series(cancel_submit_latency) * 1e-9, 'Cancel to Submit Latency')
  print('')

  oe = oe[oe['type'] == 'ORDER_SUBMITTED']
  oe['event_time'] = oe['event_time'].astype(float)
  for side, soe in oe.groupby('order_side'):
    ts = (soe['event_time'] - soe.shift(1)['event_time']) * 1e-9
    ts = ts[~np.isnan(ts)]
    output_percentile(ts, side)


if __name__ == '__main__':
  app.run(main)
